# File: app/main.py
import json
import logging
import base64

from fastapi.staticfiles import StaticFiles
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, HTTPException
from typing import Dict, Any
from datetime import datetime

from typing import Dict, Any

# message_handlers.py

from typing import Dict, Any
import logging
from datetime import datetime
from pydantic import BaseModel

from app.utils import create_jwt, decode_jwt, verify_token


app = FastAPI()

# Mount static files
app.mount("/static", StaticFiles(directory="static"), name="static")

# Mocked users and roles for demonstration
# Assuming this is your fake user database
fake_users_db: Dict[str, Dict] = {
    "user1": {"password": "password1", "role": "admin"},
    "user2": {"password": "password2", "role": "user"},
}

@app.get("/")
async def root():
    return {"message": "Go to /static/index.html to access the WebSocket login"}



class ConnectionManager:
    def __init__(self):
        self.active_connections: Dict[str, WebSocket] = {}

    async def connect(self, user_id: str, websocket: WebSocket):
        """Accepts a WebSocket connection and adds it to active connections."""
        await websocket.accept()
        self.active_connections[user_id] = websocket

    def disconnect(self, user_id: str):
        """Removes a WebSocket connection from active connections."""
        self.active_connections.pop(user_id, None)

    async def send_text(self, user_id: str, message: str):
        """Sends a text message to a specific user."""
        if user_id in self.active_connections:
            await self.active_connections[user_id].send_text(message)

    async def send_json(self, user_id: str, message: dict):
        """Sends a JSON message to a specific user."""
        if user_id in self.active_connections:
            await self.active_connections[user_id].send_json(message)

    async def send_bytes(self, user_id: str, message: bytes):
        """Sends a binary message to a specific user."""
        if user_id in self.active_connections:
            await self.active_connections[user_id].send_bytes(message)

manager = ConnectionManager()


@app.websocket("/ws/{user_id}")
async def websocket_endpoint(user_id: str, websocket: WebSocket):
    # Extract token from query parameters
    token = websocket.query_params.get("token")
    if not token:
        await websocket.close(code=1008)  # Close with error code if token is missing
        raise HTTPException(status_code=403, detail="Missing token")

    try:
        jwt_payload = verify_token(token)  # Use the new verify_token function
        logging.info(f"User {user_id} authenticated with token: {jwt_payload}")
    except HTTPException as e:
        await websocket.close(code=1008)  # Close with error code on invalid token
        raise e

    # If token is valid, connect the WebSocket
    await manager.connect(user_id, websocket)
    try:
        while True:
            # Receive ASGI websocket messages, ensuring valid state transitions
            message = await websocket.receive()
            message_type = message["type"]

            if message_type == "websocket.connect":
                logging.info(f"WebSocket connection established for user {user_id}")
                continue  # Skip to next iteration to wait for actual data

            if message_type == "websocket.receive":
                data = message.get('text', message.get('bytes', None))
                if data is None:
                    logging.warning(f"Received empty message from user {user_id}")
                    continue

                try:
                    # Try to parse as JSON
                    parsed_data = json.loads(data)
                    router_key = parsed_data.get('router')
                    if router_key:
                        # Route the message to the appropriate function
                        await route_message(user_id, router_key, parsed_data.get('data', {}))
                    else:
                        # Handle messages without a router key or log an error
                        logging.warning(f"Received JSON message without router key from user {user_id}")
                        await websocket.send_json({"status": "error", "message": "Missing router key"})
                except json.JSONDecodeError:
                    # If JSON parsing fails, treat it as plain text or binary data
                    if isinstance(data, bytes):
                        await handle_binary_data(user_id, data)
                    else:
                        await handle_text_data(user_id, data)
            elif message_type == "websocket.disconnect":
                logging.info(f"WebSocket disconnected for user {user_id}")
                manager.disconnect(user_id)
                break  # Exit the loop since the connection is closed
            else:
                # Log unexpected message types
                logging.warning(f"Received unexpected message type: {message_type} from user {user_id}")
                # Do not attempt to send a message after disconnect

    except WebSocketDisconnect:
        logging.info(f"WebSocketDisconnect for user {user_id}")
        manager.disconnect(user_id)
    except Exception as e:
        logging.error(f"WebSocket Error for user {user_id}: {str(e)}")
        await websocket.close(code=1011)


async def route_message(user_id: str, router_key: str, data: Dict[str, Any]):
    # Define a routing dictionary where keys are router keys and values are functions
    routes = {
        "function1": handle_function1,
        "function2": handle_function2,
        # Add more functions as needed
    }

    if router_key in routes:
        await routes[router_key](user_id, data)
    else:
        logging.warning(f"No handler found for router key: {router_key}")
        await manager.send_json(user_id, {"status": "error", "message": f"Unknown router key: {router_key}"})



async def handle_text_data(user_id: str, data: str):
    # Log the received text data
    logging.info(f"Received text data from user {user_id}: {data}")

    # Example: Echo the text back to the client
    response = {
        "status": "success",
        "message": "Text message received",
        "data": {
            "echo": data
        },
        "timestamp": datetime.utcnow().isoformat() + 'Z',
        "router": "handle_text_data"
    }
    await manager.send_json(user_id, response)


async def handle_binary_data(user_id: str, data: bytes):
    # Log the received binary data size
    logging.info(f"Received binary data from user {user_id}, size: {len(data)} bytes")

    # Example: Convert binary data to base64 string for easier handling or transmission
    base64_data = base64.b64encode(data).decode('utf-8')

    response = {
        "status": "success",
        "message": "Binary data received",
        "data": {
            "base64_data": base64_data
        },
        "timestamp": datetime.utcnow().isoformat() + 'Z',
        "router": "handle_binary_data"
    }
    await manager.send_json(user_id, response)





async def handle_function1(user_id: str, data: Dict[str, Any]):
    # Log the received data
    logging.info(f"User {user_id} called function1 with data: {data}")

    # Example: Process the data
    if 'message' in data:
        message = data['message']
        # You might want to do something with the message, like:
        # - Store it in a database
        # - Send it to other users
        # - Perform some analysis on it
        response = {
            "status": "success",
            "message": f"Function1 processed your message: {message}",
            "data": {
                "processed_message": message.upper()  # Just an example of processing
            },
            "timestamp": datetime.utcnow().isoformat() + 'Z',
            "router": "function1"
        }
        await manager.send_json(user_id, response)
    else:
        # If the required data is missing, send an error response
        error_response = {
            "status": "error",
            "message": "Missing 'message' in data for function1",
            "data": {},
            "timestamp": datetime.utcnow().isoformat() + 'Z',
            "router": "function1"
        }
        await manager.send_json(user_id, error_response)



async def handle_function2(user_id: str, data: Dict[str, Any]):
    # Log the received data
    logging.info(f"User {user_id} called function2 with data: {data}")

    # Example: Perform some operation based on the data
    if 'number' in data:
        number = data['number']
        try:
            # Let's say function2 is meant to square a number
            result = number ** 2
            response = {
                "status": "success",
                "message": f"Function2 squared your number: {number}",
                "data": {
                    "result": result
                },
                "timestamp": datetime.utcnow().isoformat() + 'Z',
                "router": "function2"
            }
            await manager.send_json(user_id, response)
        except TypeError:
            # If the number isn't actually a number, send an error
            error_response = {
                "status": "error",
                "message": "Invalid number provided for function2",
                "data": {},
                "timestamp": datetime.utcnow().isoformat() + 'Z',
                "router": "function2"
            }
            await manager.send_json(user_id, error_response)
    else:
        # If the required data is missing, send an error response
        error_response = {
            "status": "error",
            "message": "Missing 'number' in data for function2",
            "data": {},
            "timestamp": datetime.utcnow().isoformat() + 'Z',
            "router": "function2"
        }
        await manager.send_json(user_id, error_response)


# Pydantic model for the request body
class TokenRequest(BaseModel):
    password: str

# Example route to issue a JWT for testing

@app.post("/token/{user_id}")
async def generate_token(user_id: str, request: TokenRequest):
    # Check if the user exists in the database
    if user_id in fake_users_db:
        user_data = fake_users_db[user_id]
        # Verify the password
        if request.password == user_data["password"]:
            # If password matches, generate and return the token
            token = create_jwt({"sub": user_id, "role": user_data["role"]})
            return {"access_token": token}
        else:
            # If password does not match, raise an authentication error
            raise HTTPException(status_code=401, detail="Incorrect password")
    else:
        # If user does not exist, raise an error
        raise HTTPException(status_code=401, detail="User not found")